package com.duoec.dw.lineage.service.analyser.statement;

import com.alibaba.druid.sql.ast.SQLExpr;
import com.alibaba.druid.sql.ast.SQLStatement;
import com.alibaba.druid.sql.ast.expr.SQLIdentifierExpr;
import com.alibaba.druid.sql.ast.expr.SQLPropertyExpr;
import com.alibaba.druid.sql.ast.statement.SQLColumnDefinition;
import com.alibaba.druid.sql.ast.statement.SQLSelect;
import com.alibaba.druid.sql.ast.statement.SQLTableElement;
import com.alibaba.druid.sql.dialect.hive.stmt.HiveCreateTableStatement;
import com.duoec.dw.lineage.dto.TableFieldLineageDto;
import com.duoec.dw.lineage.dto.TableLineageDto;
import com.duoec.dw.lineage.enums.StatementType;
import com.duoec.dw.lineage.exceptions.LineageAnalyseException;
import com.duoec.dw.lineage.core.utils.AnalyserUtils;
import com.google.common.collect.Sets;
import org.springframework.stereotype.Service;
import org.springframework.util.CollectionUtils;

import java.util.LinkedHashSet;
import java.util.List;
import java.util.Set;

/**
 * @author xuwenzhen
 */
@Service
public class CreateTableStatementAnalyser extends BaseStatementAnalyser<HiveCreateTableStatement> {
    private static final int TABLE_ID_SIZE = 2;

    /**
     * 匹配具体的SQLStatement
     *
     * @return 匹配类型
     */
    @Override
    public Set<Class<? extends SQLStatement>> parseSqlStatements() {
        return Sets.newHashSet(HiveCreateTableStatement.class);
    }

    /**
     * 分析
     *
     * @param statement SQLStatement
     * @return 当前操作的血缘关系
     */
    @Override
    public TableLineageDto analyse(HiveCreateTableStatement statement) {
        TableLineageDto dto = new TableLineageDto();
        dto.setStatementType(StatementType.CREATE_TABLE);
        setTableName(dto, statement);
        setTableFields(dto, statement);
        setPartitionColumns(dto, statement);
        SQLExpr comment = statement.getComment();
        if (comment != null) {
            dto.setComment(comment.toString());
        }
        AnalyserUtils.fixFieldInfo(dto);
        return dto;
    }

    private void setPartitionColumns(TableLineageDto dto, HiveCreateTableStatement statement) {
        List<SQLColumnDefinition> partitions = statement.getPartitionColumns();
        if (CollectionUtils.isEmpty(partitions)) {
            return;
        }
        LinkedHashSet<TableFieldLineageDto> fields = dto.getFields();
        partitions.forEach(partition -> addField(fields, partition));
    }

    private void addField(LinkedHashSet<TableFieldLineageDto> fields, SQLColumnDefinition partition) {
        TableFieldLineageDto field = new TableFieldLineageDto();
        field.setColName(AnalyserUtils.formatFieldName(partition.getColumnName()));
        field.setDataType(partition.getDataType().toString());
        SQLExpr comment = partition.getComment();
        if (comment != null) {
            field.setComment(comment.toString());
        }
        field.setIndex(fields.size());
        fields.add(field);
    }

    private void setTableFields(TableLineageDto dto, HiveCreateTableStatement statement) {
        List<SQLTableElement> tableElemens = statement.getTableElementList();
        if (!CollectionUtils.isEmpty(tableElemens)) {
            setTableFieldsFromSchema(dto, tableElemens);
            return;
        }
        SQLSelect select = statement.getSelect();
        if (select != null) {
            setTableFieldsFromSelect(dto, select);
        }
    }

    private void setTableFieldsFromSelect(TableLineageDto dto, SQLSelect select) {
        TableLineageDto lineage = hiveAnalyser.analyseQuery(select.getQuery(), null);
        LinkedHashSet<TableFieldLineageDto> fields = lineage.getFields();
        dto.setFields(fields);
    }

    private void setTableFieldsFromSchema(TableLineageDto dto, List<SQLTableElement> tableElements) {
        LinkedHashSet<TableFieldLineageDto> fields = new LinkedHashSet<>();
        dto.setFields(fields);
        tableElements.forEach(element -> {
            if (element instanceof SQLColumnDefinition) {
                addField(fields, (SQLColumnDefinition) element);
            }
        });
    }

    private void setTableName(TableLineageDto dto, HiveCreateTableStatement statement) {
        SQLExpr expr = statement.getTableSource().getExpr();
        if (expr instanceof SQLPropertyExpr) {
            SQLPropertyExpr propertyExpr = (SQLPropertyExpr) expr;
            dto.setDatabaseName(AnalyserUtils.formatFieldName(propertyExpr.getOwnernName()));
            dto.setTableName(AnalyserUtils.formatFieldName(propertyExpr.getName()));
        } else if (expr instanceof SQLIdentifierExpr) {
            SQLIdentifierExpr identifierExpr = (SQLIdentifierExpr) expr;
            String name = AnalyserUtils.formatFieldName(identifierExpr.getName());
            String[] names = name.split("\\.");
            if (names.length == TABLE_ID_SIZE) {
                dto.setDatabaseName(AnalyserUtils.formatFieldName(names[0]));
                dto.setTableName(AnalyserUtils.formatFieldName(names[1]));
            } else {
                throw new LineageAnalyseException("表创建必须指定库.表，当前是：" + name);
            }
        } else {
            throw new LineageAnalyseException("未支持的表名类型：" + expr.getClass().getName() + "，当前是：" + expr.toString());
        }

    }
}
